import argparse
from datetime import time

import pandas as pd
import numpy as np
import torch
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import time
import matplotlib.pyplot as plt
import random
from tqdm import tqdm
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def set_random_seeds(seed_value):
    random.seed(seed_value)
    np.random.seed(seed_value)
    torch.manual_seed(seed_value)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed_value)
        torch.cuda.manual_seed_all(seed_value)  # 为所有GPU设置种子


# 使用示例
set_random_seeds(99)


class StandardScaler():
    def __init__(self):
        self.mean = 0.
        self.std = 1.

    def fit(self, data):
        # 计算数据的均值，并将其存储在 self.mean 中。0 参数表示沿着第0轴（通常是列）计算均值。
        self.mean = data.mean(0)
        self.std = data.std(0)

    def transform(self, data):
        mean = torch.from_numpy(self.mean).type_as(data).to(data.device) if torch.is_tensor(data) else self.mean
        std = torch.from_numpy(self.std).type_as(data).to(data.device) if torch.is_tensor(data) else self.std
        return (data - mean) / std

    def inverse_transform(self, data):
        mean = torch.from_numpy(self.mean).type_as(data).to(data.device) if torch.is_tensor(data) else self.mean
        std = torch.from_numpy(self.std).type_as(data).to(data.device) if torch.is_tensor(data) else self.std
        if data.shape[-1] != mean.shape[-1]:
            mean = mean[-1:]
            std = std[-1:]
        return (data * std) + mean


class TimeSeriesDataset(Dataset):
    def __init__(self, sequences):
        self.sequences = sequences

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, index):
        sequence, label = self.sequences[index]
        return torch.Tensor(sequence), torch.Tensor(label)


def create_inout_sequences(input_data, tw, pre_len, config):
    # 创建时间序列数据专用的数据分割器
    inout_seq = []
    L = len(input_data)
    for i in range(L - tw):
        if (i + tw + pre_len) > len(input_data):
            break
        if config.feature == 'MS':
            # train_seq = input_data[i:i + tw]
            train_seq = input_data[:, :-1][i:i + tw]
            train_label = input_data[:, -1:][i + tw:i + tw + pre_len]
        elif config.feature == 'S':
            train_seq = input_data[:, -1:][i:i + tw]
            train_label = input_data[:, -1:][i + tw:i + tw + pre_len]
        else:
            train_seq = input_data[i:i + tw]
            train_label = input_data[i + tw:i + tw + pre_len]
        inout_seq.append((train_seq, train_label))
    return inout_seq


def calculate_mae(y_true, y_pred):
    mae = np.mean(np.abs(y_true - y_pred))
    return mae


def calculate_mse(y_true, y_pred):
    mse = np.mean((y_true - y_pred) ** 2)
    return mse


def create_dataloader(config, device):
    print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>创建数据加载器<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<")
    df = pd.read_csv(config.data_path)
    pre_len = config.pre_len
    train_window = config.window_size

    # 将特征列移到末尾
    target_data = df[[config.target]]
    # 从 df DataFrame 中删除目标特征列。axis=1 表示操作的是列而不是行
    df = df.drop(config.target, axis=1)
    # 使用 pd.concat 函数将 target_data（目标特征列）重新添加到 df DataFrame 的末尾。axis=1 表示沿着列的方向进行连接
    df = pd.concat((df, target_data), axis=1)

    # 获取 df DataFrame 的列名，并存储到 cols_data 中。这里使用 [1:] 切片操作是为了排除第一列
    cols_data = df.columns[1:]
    # 根据 cols_data 中的列名从 df DataFrame 中选择对应的列，存储到 df_data 中。这样 df_data 就包含了除了第一列以外的所有特征列。
    df_data = df[cols_data]

    # 使用 .values 属性将 df_data DataFrame 中的数据转换为一个 NumPy 数组
    true_data = df_data.values

    # 定义标准化优化器
    scaler = StandardScaler()
    # 这行代码使 scaler 根据 true_data 数组中的数据计算均值和标准差
    scaler.fit(true_data)

    # 将数据集分割为训练集、验证集和测试集。
    train_data = true_data[int(0.2 * len(true_data)):]
    # valid_data = true_data[int(0.15 * len(true_data)):int(0.30 * len(true_data))]
    test_data = true_data[:int(0.2 * len(true_data))]
    print("训练集尺寸:", len(train_data), "测试集尺寸:", len(test_data))

    # 进行标准化处理
    train_data_normalized = scaler.transform(train_data)
    test_data_normalized = scaler.transform(test_data)
    # valid_data_normalized = scaler.transform(valid_data)

    # 转化为深度学习模型需要的类型Tensor
    # 将数组转换为 PyTorch 的 FloatTensor（浮点张量），然后使用 .to(device) 方法将其移动到指定的设备（GPU）
    # 因为 PyTorch 模型通常在 GPU 上运行以加速计算，而 PyTorch 张量需要在 GPU 上才能被模型处理
    train_data_normalized = torch.FloatTensor(train_data_normalized).to(device)
    test_data_normalized = torch.FloatTensor(test_data_normalized).to(device)
    # valid_data_normalized = torch.FloatTensor(valid_data_normalized).to(device)

    # 定义训练器的的输入
    # 调用 create_inout_sequences 函数，为训练数据创建输入输出序列
    # 原始的连续数据被转换成了一系列适合于训练、测试和验证深度学习模型的输入输出对
    train_inout_seq = create_inout_sequences(train_data_normalized, train_window, pre_len, config)
    test_inout_seq = create_inout_sequences(test_data_normalized, train_window, pre_len, config)
    # valid_inout_seq = create_inout_sequences(valid_data_normalized, train_window, pre_len, config)

    # 创建数据集
    train_dataset = TimeSeriesDataset(train_inout_seq)
    test_dataset = TimeSeriesDataset(test_inout_seq)
    # valid_dataset = TimeSeriesDataset(valid_inout_seq)

    # 创建 DataLoader 数据加载器，用于在模型训练和评估过程中批量加载数据。
    # batch_size 指定了每个批次的样本数量，shuffle=True 表示在每个epoch开始时随机打乱数据；
    # drop_last=True 表示如果数据集的大小不能被批次大小整除，则丢弃最后不完整的批次
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, drop_last=True)
    # valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, drop_last=True)

    print("通过滑动窗口共有训练集数据：", len(train_inout_seq), "转化为批次数据:", len(train_loader))
    print("通过滑动窗口共有测试集数据：", len(test_inout_seq), "转化为批次数据:", len(test_loader))
    # print("通过滑动窗口共有验证集数据：", len(valid_inout_seq), "转化为批次数据:", len(valid_loader))
    print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>创建数据加载器完成<<<<<<<<<<<<<<<<<<<<<<<<<<<")
    # return train_loader, test_loader, valid_loader, scaler
    return train_loader, test_loader, scaler


class GRU(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim, pre_len):
        super(GRU, self).__init__()
        self.pre_len = pre_len
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        # 替换 LSTM 为 GRU
        self.gru = nn.GRU(input_dim, hidden_dim, num_layers=num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)
        # self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        h0_gru = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)
        out, _ = self.gru(x, h0_gru)
        out = self.dropout(out)
        # 取最后 pre_len 时间步的输出
        out = out[:, -self.pre_len:, :]
        out = self.fc(out)
        # out = self.relu(out)
        return out


def train(model, args, train_loader):
    start_time = time.time()  # 计算起始时间
    model = model
    loss_function = nn.MSELoss()
    # 定义优化器
    optimizer = torch.optim.Adam(model.parameters(), lr = args.lr)
    # 从传入的 args 对象中获取训练的轮数（epochs）
    epochs = args.epochs
    model.train()  # 训练模式
    results_loss = []
    # 外层循环遍历每个 epoch。
    for i in tqdm(range(epochs)):
        losses = []
        # 内层循环遍历 train_loader 中的每个批次数据（序列和标签）
        for seq, labels in train_loader:
            optimizer.zero_grad()
            y_pred = model(seq)
            single_loss = loss_function(y_pred, labels)
            single_loss.backward()
            optimizer.step()
            losses.append(single_loss.detach().cpu().numpy())
        tqdm.write(f"\t Epoch {i + 1} / {epochs}, Loss: {sum(losses) / len(losses)}")
        results_loss.append(sum(losses) / len(losses))
        torch.save(model.state_dict(), 'save_model.pth')
        time.sleep(0.1)
    print(f">>>>>>>>>>>>>>>>>>>>>>模型已保存,用时:{(time.time() - start_time) / 60:.4f} min<<<<<<<<<<<<<<<<<<")
    # 绘制损失数据图表
    plt.plot(results_loss, label='Loss')
    # 添加标题和图例
    plt.title("GRU train loss")
    plt.legend()
    plt.show()

# 用于在验证集上评估一个已经训练好的 PyTorch 模型的性能。一个标准的验证过程，包括加载模型、设置评估模式、遍历验证集、计算损失和输出平均损失
def valid(model, args, valid_loader):
    lstm_model = model
    # 加载模型进行预测
    lstm_model.load_state_dict(torch.load('save_model.pth'))
    lstm_model.eval()  # 评估模式
    mae_losses = []
    mse_losses = []

    for seq, labels in valid_loader:
        pred = lstm_model(seq)
        mae = calculate_mae(pred.detach().cpu().numpy(), np.array(labels.detach().cpu()))  # MAE误差计算绝对值(预测值  - 真实值)
        mse = calculate_mse(pred.detach().cpu().numpy(), np.array(labels.detach().cpu()))
        mae_losses.append(mae)
        mse_losses.append(mse)

    print("验证集误差MAE:", mae_losses)
    print("验证集误差MSE:", mse_losses)
    mae = sum(mae_losses) / len(mae_losses)
    mse = sum(mse_losses) / len(mse_losses)
    print("验证集平均MAE:", mae)
    print("验证集平均MSE:", mse)


def test(model, args, test_loader, scaler):
    # 加载模型进行预测
    mae_losses = []
    mse_losses = []
    model = model
    model.load_state_dict(torch.load('save_model.pth'))
    model.eval()  # 评估模式
    results = []
    labels = []
    for seq, label in test_loader:
        pred = model(seq)
        mae = calculate_mae(pred.detach().cpu().numpy(), np.array(label.detach().cpu()))  # MAE误差计算绝对值(预测值  - 真实值)
        mse = calculate_mse(pred.detach().cpu().numpy(), np.array(label.detach().cpu()))
        mae_losses.append(mae)
        mse_losses.append(mse)
        pred = pred[:, 0, :]
        label = label[:, 0, :]
        pred = scaler.inverse_transform(pred.detach().cpu().numpy())
        label = scaler.inverse_transform(label.detach().cpu().numpy())
        for i in range(len(pred)):
            results.append(pred[i][-1])
            labels.append(label[i][-1])

    print("测试集MAE：", mae_losses)
    print("测试集MSE：", mse_losses)
    average_mae = np.mean(mae_losses)
    average_mse = np.mean(mse_losses)
    print("测试集平均误差MAE:", average_mae)
    print("测试集均方误差MSE:", average_mse)
    # 绘制历史数据
    plt.plot(labels, label='TrueValue')
    # 绘制预测数据
    # 注意这里预测数据的起始x坐标是历史数据的最后一个点的x坐标
    plt.plot(results, label='Prediction')
    # 添加标题和图例
    plt.title("GRU test state")
    plt.legend()
    plt.show()

# 检验模型拟合情况，提供了一个直观的方式来检查模型在训练集上的拟合情况，通过绘制预测结果与真实值的对比图
def inspect_model_fit(model, args, train_loader, scaler):
    model = model
    model.load_state_dict(torch.load('save_model.pth'))
    model.eval()  # 评估模式
    results = []
    labels = []

    for seq, label in train_loader:
        pred = model(seq)[:, 0, :]
        label = label[:, 0, :]
        pred = scaler.inverse_transform(pred.detach().cpu().numpy())
        label = scaler.inverse_transform(label.detach().cpu().numpy())
        for i in range(len(pred)):
            results.append(pred[i][-1])
            labels.append(label[i][-1])

    # 绘制历史数据
    plt.plot(labels, label='History')

    # 绘制预测数据
    # 注意这里预测数据的起始x坐标是历史数据的最后一个点的x坐标
    plt.plot(results, label='Prediction')

    # 添加标题和图例
    plt.title("GRU inspect model fit state")
    plt.legend()
    plt.show()


def predict(model, args, device):
    # 预测未知数据的功能
    # 从指定路径读取 CSV 文件到 pandas DataFrame
    df = pd.read_csv(args.data_path)
    if args.feature == 'S':
        # 提取 DataFrame 的最后 args.window_size 行，并转换为 NumPy 数组
        df = df.iloc[:, -1:][-args.window_size:].values  # 转换为nadarry
    elif args.feature == 'MS':
        # 提取 DataFrame 的最后 args.window_size 行，并转换为 NumPy 数组
        df = df.iloc[:, 1:-1][-args.window_size:].values  # 转换为nadarry
    # 定义标准化优化器
    scaler = StandardScaler()
    scaler.fit(df)
    # 使用 scaler 对特征数据进行标准化处理
    pre_data = scaler.transform(df)
    # 将 NumPy 数组转换为 PyTorch 张量，并移动到指定的设备（如 GPU）
    tensor_pred = torch.FloatTensor(pre_data).to(device)
    # 增加一个批次维度，因为 PyTorch 模型通常期望批次维度
    tensor_pred = tensor_pred.unsqueeze(0)  # 单次预测 , 滚动预测功能暂未开发后期补上
    model = model
    model.load_state_dict(torch.load('save_model.pth'))
    model.eval()  # 评估模式

    # 使用模型对处理后的数据进行预测
    pred = model(tensor_pred)[0]

    # 使用 scaler 对预测结果进行逆变换，恢复到原始尺度
    pred = scaler.inverse_transform(pred.detach().cpu().numpy())

    # 假设 df 和 pred 是你的历史和预测数据

    # 计算历史数据的长度
    history_length = len(df[:, -1])

    # 为历史数据生成x轴坐标
    history_x = range(history_length)

    # 为预测数据生成x轴坐标
    # 开始于历史数据的最后一个点的x坐标
    prediction_x = range(history_length - 1, history_length + len(pred[:, -1]) - 1)

    # 绘制历史数据
    plt.plot(history_x, df[:, -1], label='History')

    # 绘制预测数据
    # 注意这里预测数据的起始x坐标是历史数据的最后一个点的x坐标
    plt.plot(prediction_x, pred[:, -1], label='Prediction')
    plt.axvline(history_length - 1, color='red')  # 在图像的x位置处画一条红色竖线
    # 添加标题和图例
    plt.title("GRU History and Prediction")
    plt.legend()


if __name__ == '__main__':
    # 创建一个 ArgumentParser 对象，用于处理命令行参数，描述信息设置为 "Time Series forecast"
    parser = argparse.ArgumentParser(description='Time Series forecast')
    parser.add_argument('--model', type=str, default='LSTM', help="模型持续更新")
    parser.add_argument('--window_size', type=int, default=23, help="时间窗口大小, window_size > pre_len")
    parser.add_argument('--pre_len', type=int, default=12, help="预测未来数据长度")
    # old_data
    parser.add_argument('--shuffle', action='store_true', default=True, help="是否打乱数据加载器中的数据顺序")
    parser.add_argument('--data_path', type=str, default='../old_data/all_add.csv', help="你的数据数据地址")
    parser.add_argument('--target', type=str, default='RSEI', help='你需要预测的特征列，这个值会最后保存在csv文件里')
    parser.add_argument('--input_size', type=int, default=1, help='你的特征个数不算时间和目标列的数量')
    parser.add_argument('--feature', type=str, default='S', help='[M, S, MS],多元预测多元,单元预测单元,多元预测单元')
    parser.add_argument('--model_dim', type=list, default=64, help='隐藏层神经元数量')

    # learning
    parser.add_argument('--lr', type=float, default=0.002, help="学习率")
    parser.add_argument('--drop_out', type=float, default=0.3, help="随机丢弃概率,防止过拟合")
    parser.add_argument('--epochs', type=int, default=150, help="训练轮次")
    parser.add_argument('--batch_size', type=int, default=4, help="批次大小")
    parser.add_argument('--save_path', type=str, default='models')

    # model
    parser.add_argument('--hidden_size', type=int, default=64, help="隐藏层单元数")
    parser.add_argument('--kernel_sizes', type=int, default=3)
    parser.add_argument('--laryer_num', type=int, default=1)
    # device
    parser.add_argument('--use_gpu', type=bool, default=True)
    parser.add_argument('--device', type=int, default=0, help="只设置最多支持单个gpu训练")

    # option
    parser.add_argument('--train', type=bool, default=True)
    parser.add_argument('--valid', type=bool, default=True)
    parser.add_argument('--test', type=bool, default=True)
    parser.add_argument('--predict', type=bool, default=True)
    parser.add_argument('--inspect_fit', type=bool, default=True)
    parser.add_argument('--lr-scheduler', type=bool, default=True)
    # 解析命令行输入的参数，并将它们存储在 args 对象中
    args = parser.parse_args()

    if isinstance(args.device, int) and args.use_gpu:
        device = torch.device("cuda:" + f'{args.device}')
    else:
        device = torch.device("cpu")
    print("使用设备:", device)
    # 调用 create_dataloader 函数，根据 args 和 device 创建训练、测试和验证的数据加载器 (DataLoader) 以及数据标准化器 (scaler)
    # train_loader, test_loader, valid_loader, scaler = create_dataloader(args, device)
    train_loader, test_loader, scaler = create_dataloader(args, device)

    if args.feature == 'MS' or args.feature == 'S':
        args.output_size = 1
    else:
        args.output_size = args.input_size

    # 实例化模型
    try:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始初始化{args.model}模型<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        model = GRU(args.input_size, args.model_dim, args.laryer_num, args.output_size, args.pre_len).to(device)
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始初始化{args.model}模型成功<<<<<<<<<<<<<<<<<<<<<<<<<<<")
    except:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始初始化{args.model}模型失败<<<<<<<<<<<<<<<<<<<<<<<<<<<")

    # 训练模型
    if args.train:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始{args.model}模型训练<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        train(model, args, train_loader)
    # if args.valid:
    #     print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始{args.model}模型训练<<<<<<<<<<<<<<<<<<<<<<<<<<<")
    #     valid(model, args, valid_loader)
    if args.test:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始{args.model}模型测试<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        test(model, args, test_loader, scaler)
    if args.inspect_fit:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>开始检验{args.model}模型拟合情况<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        inspect_model_fit(model, args, train_loader, scaler)
    if args.predict:
        print(f">>>>>>>>>>>>>>>>>>>>>>>>>预测未来{args.pre_len}条数据<<<<<<<<<<<<<<<<<<<<<<<<<<<")
        predict(model, args, device)
    plt.show()